import matplotlib.pyplot as plt
import numpy as np



# =====================
plt.rcParams.update({
    'font.family': 'Times New Roman',
    'font.size': 12,
    'axes.labelsize': 14,
    'axes.titlesize': 16,
    'xtick.labelsize': 12,
    'ytick.labelsize': 12,
    'legend.fontsize': 11,
    'figure.figsize': (10, 6),
    'figure.dpi': 300,
    'axes.grid': True,
    'grid.linestyle': '--',
    'grid.alpha': 0.4
})

# =====================
# 实验数据
# =====================
clients = [10, 30, 50, 100]  # 客户端数量
fedavg_comm = [320, 680, 1120, 2250]  # FedAvg单轮通信量(KB)
fedatt_comm = [280, 580, 950, 1900]   # FedAtt通信量
fedent_comm = [75, 165, 280, 550]     # FedEntGate压缩后通信量
compression_ratio = [fedent_comm[i]/fedavg_comm[i] for i in range(len(clients))]  # 压缩率

# =====================
# 创建双Y轴图表
# =====================
fig, ax1 = plt.subplots()

# 左侧纵轴：通信数据量柱状图
bar_width = 0.25
x = np.arange(len(clients))
ax1.bar(x - bar_width, fedavg_comm, width=bar_width, 
        label='FedAvg', color='#1f77b4', alpha=0.8)
ax1.bar(x, fedatt_comm, width=bar_width, 
        label='FedAtt', color='#ff7f0e', alpha=0.8)
ax1.bar(x + bar_width, fedent_comm, width=bar_width, 
        label='FedEntGate', color='#2ca02c', alpha=0.8)

# 添加数据标签
for i, comm in enumerate(fedent_comm):
    ax1.text(i + bar_width, comm + 50, f'{comm}KB', 
             ha='center', fontsize=10, fontweight='bold', color='#2ca02c')

ax1.set_xticks(x)
ax1.set_xticklabels([f'{c} Clients' for c in clients])
ax1.set_xlabel('Number of Clients', fontweight='bold')
ax1.set_ylabel('Communication per Round (KB)', color='black', fontweight='bold')
ax1.tick_params(axis='y', labelcolor='black')
ax1.set_ylim(0, 2500)

# 右侧纵轴：压缩率折线图
ax2 = ax1.twinx()
ax2.plot(x, compression_ratio, 'ko--', linewidth=2, markersize=8, 
         label='Compression Ratio')
ax2.set_ylabel('Compression Ratio (FedEntGate/FedAvg)', color='black', fontweight='bold')
ax2.tick_params(axis='y', labelcolor='black')
ax2.set_ylim(0, 0.35)
ax2.yaxis.set_major_formatter('{x:.2f}×')

# 添加压缩率标注
# for i, ratio in enumerate(compression_ratio):
#     ax2.text(i, ratio + 0.02, f'{ratio:.2f}×', 
#              ha='center', fontsize=10, fontweight='bold')

# =====================
# 图表装饰与说明
# =====================
ax1.legend(loc='upper left', frameon=True, edgecolor='black')
ax2.legend(loc='upper right', frameon=True, edgecolor='black')

# plt.title('(a) Communication Efficiency Optimization of FedEntGate\n'
#          'Structured Pruning + Quantization Encoding Reduce 76% Traffic', 
#          fontsize=15, pad=20)

# 添加技术说明
# plt.annotate('Compression Mechanism:\n'
#              '- Structured Pruning: Remove 48% redundant parameters\n'
#              '- 8-bit Quantization: 4× volume reduction\n'
#              '- Huffman Encoding: Additional 20% compression', 
#              xy=(0.5, 0.75), xycoords='axes fraction',
#              bbox=dict(boxstyle="round,pad=0.5", fc="lavender", ec="b", lw=1))

plt.tight_layout()
plt.savefig('Fig6 Communication_Efficiency.png', dpi=300, bbox_inches='tight')
plt.show()
